"""
计算函数
"""
import numpy as np
import xarray as xr
from typing import Literal
from ..physics.constant import LENGTH_PRE_LATITUDE as _LENGTH_PRE_LATITUDE
from pint import Quantity

def lon_lat_distance(lon: np.ndarray, lat: np.ndarray
                     ) -> tuple[ Quantity, Quantity ]:
    """根据等间距经纬度，生成实际经纬网格间距。使用球形模型，地球半径来自WGS84.

    Generate the actual grid spacing of latitude and longitude based on latitude and longitude.
    The spherical model is used, and the radius of the earth comes from WGS84.

    Parameters
    ----------
    lon : np.ndarray
        经度  longitude
    lat : np.ndarray
        纬度  latitude 

    Returns
    -------
    tuple[ Quantity, Quantity ]
        返回纬向（x）和经向（y）的网格间距数组，单位为米

        return the grid spacing array of latitude (x) and longitude (y), in meters
    """

    # 验证输入是否为1维
    if lon.ndim != 1 or lat.ndim != 1:
        raise ValueError("lon and lat must be 1D arrays")

    # 验证是否等间距
    if not np.all(np.isclose(np.diff(lon), np.diff(lon)[0])):
        raise ValueError("Longitude is not equidistant")
    if not np.all(np.isclose(np.diff(lat), np.diff(lat)[0])):
        raise ValueError("Latitude is not equidistant")

    resolutionX: float = np.abs(lon[1] - lon[0])
    resolutionY: float = np.abs(lat[1] - lat[0])

    dy = np.ones((len(lat), len(lon))) * _LENGTH_PRE_LATITUDE * resolutionY
    dx = np.ones((len(lat))) * _LENGTH_PRE_LATITUDE * \
        np.cos(lat * np.pi / 180) * resolutionX
    dx = np.expand_dims(dx, axis=1).repeat(len(lon), axis=1)
    return dx, dy


def derivative(data: xr.DataArray, coord: str,
                dim_type: Literal["auto", "horizontal", "other"] = "auto",
                lon_name: str = "longitude", lat_name: str = "latitude",
                datetime_unit: str = None, edge_order: Literal[1, 2] = 2
                ) -> xr.DataArray:
    """计算物理量在指定坐标上的一阶导数/梯度

    Parameters
    ----------
    data : xr.DataArray
        要计算的数据，目前仅支持 `DataArray` 类型

        The data to be calculated, currently only supports `DataArray` type
    coord : str
        要求导数的维度名称

        The name of the dimension to be differentiated

    dim_type : Literal["auto", "horizontal", "other"], optional
        维度的类型，是空间维度还是其他，默认为自动识别

        The type of dimension, whether it is a spatial dimension or other, 
        default is auto-detection.

    lon_name, lat_name : str, optional
        数据中经纬度维度的名称，默认为 "longitude" 和 "latitude"

        The name of the longitude and latitude dimensions in the data, default 
        is "longitude" and "latitude"

    datetime_unit : str, optional
        计算时间变化率时使用的时间单位。

        The time unit used when calculating the rate of change with time.

    edge_order : Literal[1, 2], optional
        计算空间导数时，使用的边界处理方式，默认为2阶

        The boundary processing method used when calculating the spatial 
        derivative, default is 2nd order.

    Returns
    -------
    xr.DataArray
        _description_
    """
    data = data.copy()
    if dim_type == "auto":
        if coord in ["lon", "lat", "longitude", "latitude", lon_name, lat_name]:
            dim_type = "horizontal"
        else:
            dim_type = "other"

    if dim_type == "horizontal":
        dx, dy = lon_lat_distance(data[lon_name].values, data[lat_name].values)
        axis = data.get_axis_num(coord)
        data.values = np.gradient(data.values, axis=axis, edge_order=edge_order)

        if coord == lon_name:
            data.values /= dx
        else:
            if data[lat_name].values[0] > data[lat_name].values[-1]:
                dy *= -1
            data.values /= dy
        
        return data
    elif dim_type == "other":
        return data.differentiate(coord, edge_order=edge_order, 
                                    datetime_unit=datetime_unit)
